from paddleocr import PaddleOCR
import numpy as np
from typing import List, Dict
import logging

class OCREngine:
    def __init__(self):
        # 初始化中英文OCR模型
        try:
            # 使用英文模型识别外文
            self.ocr = PaddleOCR(
                use_angle_cls=True,
                lang='en',
                use_gpu=False
            )
            logging.info("OCR引擎初始化成功")
        except Exception as e:
            logging.error(f"OCR引擎初始化失败: {e}")
            raise
    
    def detect_text(self, image_path: str) -> List[Dict]:
        """识别图片中的文字和位置"""
        try:
            result = self.ocr.ocr(image_path, cls=True)
            if not result or not result[0]:
                return []
            
            # 格式化识别结果
            formatted_results = []
            for detection in result[0]:
                if detection and detection[1]:
                    text, confidence = detection[1]
                    formatted_results.append({
                        "text": text,
                        "confidence": float(confidence),
                        "position": np.array(detection[0]).tolist()  # 四个角点坐标
                    })
            return formatted_results
        except Exception as e:
            logging.error(f"OCR识别失败: {e}")
            return []